package track_formats

import (
	"time"

	"github.com/pion/rtp"
	"gogs.smarteye.com/servere/gobvrtp/encoder"
	"gogs.smarteye.com/servere/gobvrtp/track"
)

type PSTrack struct {
	VideoTrack

	CodecIdList []track.CodecId

	encoder *encoder.PSEncoder
	decoder *encoder.PSDecoder
}

func NewPSTrack() *PSTrack {
	c := &PSTrack{}

	param := c.CtlParma()
	param.SetPayloadType(112)
	param.SetCodec(track.CodecIdPS)
	param.SetClockRate(90000)
	param.SetOnReadRTPFunc(c.onRTP)

	return c
}

func (c *PSTrack) AddStream(cids ...track.CodecId) *PSTrack {
	c.CodecIdList = append(c.CodecIdList, cids...)
	return c
}

func (c *PSTrack) WritePacketEx(t time.Time, data []byte, codec track.CodecId) error {
	param := c.CtlParma()

	if c.encoder == nil {
		// 初始化编码器
		e := &encoder.PSEncoder{
			RTPEncoder: encoder.RTPEncoder{
				SampleRate:            c.ClockRate(),
				PayloadType:           c.PayloadType(),
				SSRC:                  ToPtr(c.SSRC()),
				InitialSequenceNumber: ToPtr(param.InitialSequenceNumber),
				PayloadMaxSize:        param.PayloadMaxSize,
			},
			CodecIdList: c.CodecIdList,
		}
		_ = e.Init()

		c.encoder = e
	}

	packets, err := c.encoder.Encode(codec, data)
	if err != nil {
		return err
	}

	timestamp, _ := c.CtlParma().RtpTimeEncode(t)

	for _, pkt := range packets {
		packet := pkt
		packet.Timestamp = timestamp
		if err = c.WriteRTP(packet); err != nil {
			return err
		}
	}

	return nil
}

func (c *PSTrack) onRTP(_ track.Track, pkt *rtp.Packet) {
	packets, _ := c.Detector(pkt)

	if c.decoder == nil {
		// 初始化解码器
		decoder := &encoder.PSDecoder{
			RTPDecoder: encoder.RTPDecoder{
				SampleRate: c.ClockRate(),
			},
		}
		_ = decoder.Init()
		c.decoder = decoder
	}

	for _, packet := range packets {
		_ = c.decoder.Decode(packet, func(data []byte, pts time.Duration, codec track.CodecId) {
			frame := &track.Frame{
				Codec: codec,
				Data:  data,
			}
			frame.PTS, _ = c.PacketPTS(pkt)
			frame.NTP, _ = c.PacketNTP(pkt)
			frame.Codec = c.Codec()
			c.CtlParma().OnFrame(c, frame)
		})
	}
}

func init() {
	RegisterTrack(track.CodecIdPS, func() track.Track {
		return NewPSTrack()
	})
}
